#!/usr/bin/env python

# 将trf的输出文件 转化为gff3格式

import argparse
import sys

from Bio import SeqIO
import re

parser = argparse.ArgumentParser(
    description='''将gff3转化为zff，用于snap训练，由大天才于2021年6月26日创建于浙江农业大学
    03_gff32zff.py -i train_for_snap.gff -g NIG_ONT_curated_v1.2.2.fasta.masked.trf_masked -o train.zff -d genome.dna''')

parser.add_argument('-i',
                help='必须给定，输入的gff3文件')

parser.add_argument('-g',
                help='必须给定，输入的参考基因组fasta文件')

parser.add_argument('-o',
                help='必须给定，输出的zff文件')

parser.add_argument('-d',
                help='必须给定，修正排序后的fasta文件')

args = parser.parse_args()

if not args.i or not args.o or not args.g or not args.d:
    parser.print_help()
    sys.exit()




gff_file = args.i


outfile = args.o


genome_file = args.g

out_genome_file = args.d


reParent = re.compile('Parent=([^;]+)')



gene_dic = {}
with open(gff_file) as fila:
    for i in fila:
        k = i.strip().split('\t')
        if len(k)>5 and k[2]=='CDS':
            r = re.search(reParent, k[-1])
            if r:
                parent = r.groups()[0]
                if parent not in gene_dic:
                    gene_dic[parent] = [k]
                else:
                    gene_dic[parent].append(k)

scaf_dic = {}
for i in gene_dic:
    scaf = gene_dic[i][0][0]
    if scaf not in scaf_dic:
        scaf_dic[scaf] = [gene_dic[i]]
    else:
        scaf_dic[scaf].append(gene_dic[i])

with open(outfile,'w') as fila:

    for scaf in scaf_dic:
        scaf_dic[scaf] = sorted(scaf_dic[scaf],key=lambda x:int(x[0][3]))
        fila.write('>'+scaf+'\n')
        for cds in scaf_dic[scaf]:
            p = re.search(reParent, cds[0][-1]).groups()[0]
            k = cds[0]
            fila.write('Einit %s %s %s 100 0 0 %s %s\n' %(k[3],k[4],k[6],k[7],p))
            for k in cds[1:-1]:
                fila.write('Exon %s %s %s 100 0 0 %s %s\n' %(k[3],k[4],k[6],k[7],p))
            k = cds[-1]
            fila.write('Eterm %s %s %s 100 0 0 %s %s\n' %(k[3],k[4],k[6],k[7],p))

genome_dic = {}

for i in SeqIO.parse(genome_file,'fasta'):
    name = str(i.name)
    seq = str(i.seq)
    genome_dic[name] = seq


with open(out_genome_file ,'w') as fita:
    for scaf in scaf_dic:
        fita.write('>'+scaf+'\n')
        fita.write(genome_dic[scaf]+'\n')

